package com.base.data.elasticsearch.repository.query;

/*
 * Copyright 2013-2017 the original author or authors.
 *
 * Licensed under the Apache License, Version 2.0 (the "License");
 * you may not use this file except in compliance with the License.
 * You may obtain a copy of the License at
 *
 *      http://www.apache.org/licenses/LICENSE-2.0
 *
 * Unless required by applicable law or agreed to in writing, software
 * distributed under the License is distributed on an "AS IS" BASIS,
 * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 * See the License for the specific language governing permissions and
 * limitations under the License.
 */

import com.base.data.elasticsearch.core.mapping.BaseElasticsearchPersistentProperty;
import com.base.data.elasticsearch.core.operation.BaseElasticsearchOperations;
import com.base.data.elasticsearch.core.query.BaseCriteriaQuery;
import com.base.data.elasticsearch.repository.query.parse.BaseElasticsearchQueryCreator;
import org.springframework.data.domain.PageRequest;
import org.springframework.data.mapping.context.MappingContext;
import org.springframework.data.repository.query.ParametersParameterAccessor;
import org.springframework.data.repository.query.parser.PartTree;
import org.springframework.data.util.CloseableIterator;
import org.springframework.data.util.StreamUtils;
import org.springframework.util.ClassUtils;

/**
 * ElasticsearchPartQuery
 */
public class BaseElasticsearchPartQuery extends BaseAbstractElasticsearchRepositoryQuery {

    private final PartTree tree;
    private final MappingContext<?, BaseElasticsearchPersistentProperty> mappingContext;

    public BaseElasticsearchPartQuery(BaseElasticsearchQueryMethod method, BaseElasticsearchOperations elasticsearchOperations) {
        super(method, elasticsearchOperations);
        this.tree = new PartTree(method.getName(), method.getEntityInformation().getJavaType());
        this.mappingContext = elasticsearchOperations.getElasticsearchConverter().getMappingContext();
    }

    @Override
    public Object execute(Object[] parameters) {
        ParametersParameterAccessor accessor = new ParametersParameterAccessor(queryMethod.getParameters(), parameters);
        BaseCriteriaQuery query = createQuery(accessor);
        if(tree.isDelete()) {
            Object result = countOrGetDocumentsForDelete(query, accessor);
            elasticsearchOperations.delete(query, queryMethod.getEntityInformation().getJavaType());
            return result;
        } else if (queryMethod.isPageQuery()) {
            query.setPageable(accessor.getPageable());
            return elasticsearchOperations.queryForPage(query, queryMethod.getEntityInformation().getJavaType());
        } else if (queryMethod.isStreamQuery()) {
            Class<?> entityType = queryMethod.getEntityInformation().getJavaType();
            if (query.getPageable().isUnpaged()) {
                int itemCount = (int) elasticsearchOperations.count(query, queryMethod.getEntityInformation().getJavaType());
                query.setPageable(PageRequest.of(0, Math.max(1, itemCount)));
            }

            return StreamUtils.createStreamFromIterator((CloseableIterator<Object>) elasticsearchOperations.stream(query, entityType));

        } else if (queryMethod.isCollectionQuery()) {
            if (accessor.getPageable() == null) {
                int itemCount = (int) elasticsearchOperations.count(query, queryMethod.getEntityInformation().getJavaType());
                query.setPageable(PageRequest.of(0, Math.max(1, itemCount)));
            } else {
                query.setPageable(accessor.getPageable());
            }
            return elasticsearchOperations.queryForList(query, queryMethod.getEntityInformation().getJavaType());
        } else if (tree.isCountProjection()) {
            return elasticsearchOperations.count(query, queryMethod.getEntityInformation().getJavaType());
        }
        return elasticsearchOperations.queryForObject(query, queryMethod.getEntityInformation().getJavaType());
    }

    private Object countOrGetDocumentsForDelete(BaseCriteriaQuery query, ParametersParameterAccessor accessor) {

        Object result = null;

        if (queryMethod.isCollectionQuery()) {
            if (accessor.getPageable().isUnpaged()) {
                int itemCount = (int) elasticsearchOperations.count(query, queryMethod.getEntityInformation().getJavaType());
                query.setPageable(PageRequest.of(0, Math.max(1, itemCount)));
            } else {
                query.setPageable(accessor.getPageable());
            }
            result = elasticsearchOperations.queryForList(query, queryMethod.getEntityInformation().getJavaType());
        }

        if (ClassUtils.isAssignable(Number.class, queryMethod.getReturnedObjectType())) {
            result = elasticsearchOperations.count(query, queryMethod.getEntityInformation().getJavaType());
        }
        return result;
    }

    public BaseCriteriaQuery createQuery(ParametersParameterAccessor accessor) {
        return new BaseElasticsearchQueryCreator(tree, accessor, mappingContext).createQuery();
    }
}
